#include <stdio.h>
#include <cuda_runtime.h>

// 定义向量长度
#define N (16 * 1024 * 1024)

// CUDA Kernel：向量加法（每个线程处理一个元素）
__global__ void vectorAdd(const float *A, const float *B, float *C, int numElements) {
	int i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < numElements) {
		C[i] = A[i] + B[i];  // 逐元素相加
	}
}

// 主函数
int main() {
	size_t size = N * sizeof(float);

	// 使用统一内存分配空间
	float *a, *b, *c;
	cudaMallocManaged(&a, size);
	cudaMallocManaged(&b, size);
	cudaMallocManaged(&c, size);

	for (int i = 0; i < N; i++) {
		a[i] = rand() / (float)RAND_MAX;
		b[i] = rand() / (float)RAND_MAX;
	}

	// 定义线程块和网格大小
	int threadsPerBlock = 256;
	int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;

	// 创建计时事件
	cudaEvent_t start, stop;
	cudaEventCreate(&start);
	cudaEventCreate(&stop);

	//开始计时
	cudaEventRecord(start);

	cudaStream_t stream;
	cudaStreamCreate(&stream);
	cudaMemPrefetchAsync(a, size, 0, stream);
	cudaMemPrefetchAsync(b, size, 0, stream);

	// 启动CUDA Kernel
	vectorAdd<<<blocksPerGrid, threadsPerBlock, 0, stream>>>(a, b, c, N);

	// 同步流
	cudaStreamSynchronize(stream);

	//停止计时
	cudaEventRecord(stop);
	cudaEventSynchronize(stop);

	// 计算时间
	float totalMs;
	cudaEventElapsedTime(&totalMs, start, stop);
	printf("数据量: %d 个元素\n", N);
	printf("总耗时: %.3f ms\n", totalMs);

	// 验证前10个结果
	for (int i = 0; i < 10; i++) {
		printf("c[%d] = %.2f (expected: %.2f)\n", i, c[i], a[i] + b[i]);
	}

	// 释放内存
	cudaFree(a);
	cudaFree(b);
	cudaFree(c);

	//销毁Stream
	cudaStreamDestroy(stream);

	//销毁Event
	cudaEventDestroy(start);
	cudaEventDestroy(stop);

	return 0;
}
